use crate::common::{config::ServerConfig, errors::Result};
use crate::common::proto::{self, oc_cluster};
use crate::server::registry::{NODE_REGISTRY, RegisteredNode};
use crate::common::storage::Storage;
use tokio_rustls::TlsAcceptor;
use rustls::{ServerConfig as RustlsConfig, PrivateKey, Certificate};
use std::sync::Arc;
use tokio::net::TcpListener;
use tracing::{info, error, debug};
use crate::common::metrics::METRICS;

pub async fn start_tls_server(config: ServerConfig, storage: std::sync::Arc<dyn Storage>) -> Result<()> {
    let addr = format!("{}:{}", config.server.host, config.server.port);
    let listener = TcpListener::bind(&addr).await?;
    
    info!("TLS server listening on {}", addr);
    
    // 加载TLS配置
    let tls_config = load_tls_config(&config.tls)?;
    let acceptor = TlsAcceptor::from(Arc::new(tls_config));
    
    loop {
        match listener.accept().await {
            Ok((stream, addr)) => {
                debug!("New connection from {}", addr);
                let acceptor = acceptor.clone();
                
                tokio::spawn(async move {
                    if let Err(e) = handle_connection(stream, acceptor, storage).await {
                        error!("Connection error: {}", e);
                    }
                });
            }
            Err(e) => {
                error!("Accept error: {}", e);
            }
        }
    }
}

fn load_tls_config(tls_config: &crate::common::tls_config::TlsConfig) -> Result<RustlsConfig> {
    // 读取证书文件
    let cert_file = std::fs::read(&tls_config.cert_file)?;
    let key_file = std::fs::read(&tls_config.key_file)?;
    
    let certs = rustls_pemfile::certs(&mut &cert_file[..])?
        .into_iter()
        .map(Certificate)
        .collect();
    
    let key = rustls_pemfile::pkcs8_private_keys(&mut &key_file[..])?
        .first()
        .ok_or_else(|| crate::common::errors::ManagerError::CertificateError("No private key found".to_string()))?
        .clone();
    
    let mut builder = RustlsConfig::builder().with_safe_defaults();

    // 客户端认证开关
    let config = if tls_config.verify_client {
        // 若启用客户端认证，需要提供 CA 证书
        let mut roots = rustls::RootCertStore::empty();
        if let Some(ca_path) = &tls_config.ca_file {
            let ca_bytes = std::fs::read(ca_path)?;
            let cas = rustls_pemfile::certs(&mut &ca_bytes[..])?;
            for c in cas { roots.add(&rustls::Certificate(c))?; }
        }
        builder
            .with_client_cert_verifier(Arc::new(rustls::server::AllowAnyAuthenticatedClient::new(roots)))
            .with_single_cert(certs, PrivateKey(key))?
    } else {
        builder
            .with_no_client_auth()
            .with_single_cert(certs, PrivateKey(key))?
    };
    
    Ok(config)
}

async fn handle_connection(
    stream: tokio::net::TcpStream,
    acceptor: TlsAcceptor,
    storage: std::sync::Arc<dyn Storage>,
) -> Result<()> {
    let mut tls_stream = acceptor.accept(stream).await?;

    use tokio::io::{AsyncReadExt, AsyncWriteExt};
    use prost::Message;

    debug!("TLS connection established");

    // 循环读取消息，支持 Register + Heartbeat
    let mut buf = Vec::new();
    loop {
        match proto::read_frame(&mut tls_stream, &mut buf).await {
            Ok(_) => {},
            Err(e) => { error!("read error: {}", e); break; }
        }

        match oc_cluster::ClusterMessage::decode(&*buf) {
            Ok(msg) => {
                match oc_cluster::MessageType::from_i32(msg.r#type) {
                    Some(oc_cluster::MessageType::Register) => {
                        match oc_cluster::NodeInfo::decode(&*msg.payload) {
                            Ok(node_info) => {
                                let node = RegisteredNode {
                                    id: node_info.id,
                                    hostname: node_info.hostname,
                                    ip: node_info.ip_address,
                                    tags: node_info.tags,
                                    version: env!("CARGO_PKG_VERSION").to_string(),
                                    last_seen: chrono::Utc::now(),
                                };
                                NODE_REGISTRY.insert_or_update(node.clone());
                                info!("Node registered: {} {}", node.hostname, node.ip);
                                // 持久化到数据库
                                let db_node = crate::common::storage::NodeInfo {
                                    id: node.id.clone(),
                                    hostname: node.hostname.clone(),
                                    ip_address: node.ip.clone(),
                                    tags: node.tags.clone(),
                                    status: crate::common::storage::NodeStatus::Online,
                                    last_heartbeat: node.last_seen,
                                    capabilities: vec![],
                                    metadata: Default::default(),
                                    group_ids: vec![],
                                };
                                if let Err(e) = storage.register_node(&db_node).await {
                                    error!("db register_node error: {}", e);
                                }
                            }
                            Err(e) => error!("decode NodeInfo error: {}", e),
                        }
                    }
                    Some(oc_cluster::MessageType::Heartbeat) => {
                        match oc_cluster::Heartbeat::decode(&*msg.payload) {
                            Ok(hb) => {
                                let now = chrono::Utc::now();
                                NODE_REGISTRY.update_last_seen(&hb.node_id, now);
                                debug!("heartbeat received from {}", hb.node_id);
                                METRICS.record_node_heartbeat(now);
                                if let Err(e) = storage.update_node_heartbeat(&hb.node_id, now).await {
                                    error!("db update_node_heartbeat error: {}", e);
                                }
                            }
                            Err(e) => error!("decode Heartbeat error: {}", e),
                        }
                    }
                    _ => {
                        debug!("unexpected message type: {}", msg.r#type);
                    }
                }
            }
            Err(e) => error!("decode frame error: {}", e),
        }
    }

    Ok(())
} 